import torch
import torch.nn as nn
import torch.nn.functional as F
from modeling.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
from modeling.aspp import build_aspp
from modeling.decoder import build_decoder
from modeling.backbone import build_backbone
from modeling.backbone import resnet, xception, drn, mobilenet
import torchsummary


class Network(nn.Module):
	def __init__(self, output_stride=8, num_classes=1, sync_bn=True, cls_classes=3):
		super(Network, self).__init__()
		self.backone = 'resnet'
		
		if output_stride == 32:
			strides = [2, 2, 2, 2]

		elif output_stride == 16:
			strides = [2, 2, 2, 1]
		else:
			strides = [2, 2, 1, 1]
		
		if sync_bn:
			BatchNorm = SynchronizedBatchNorm2d
		else:
			BatchNorm = nn.BatchNorm2d
		
		self.backone_cls = resnet.ResNet101(output_stride, BatchNorm, pretrained=True)
		self.backone_seg = resnet.ResNet101(output_stride, BatchNorm, pretrained=True)
		
		self.aspp_s = build_aspp(self.backone, output_stride, BatchNorm)
		self.conv_cs = nn.Sequential(
			nn.Conv2d(2048, 256, kernel_size=1, stride=1, padding=0, bias=False),
			nn.BatchNorm2d(256),
			nn.ReLU(inplace=True))
		
		self.conv_s_cs = nn.Sequential(
			nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0, bias=False),
			nn.BatchNorm2d(256),
			nn.ReLU(inplace=True))
		
		self.conv_down1 = nn.Sequential(
			nn.Conv2d(256, 512, kernel_size=3, stride=strides[0], padding=1, bias=False),
			nn.BatchNorm2d(512),
			nn.ReLU(inplace=True))
		
		self.conv_down2 = nn.Sequential(
			nn.Conv2d(512, 1024, kernel_size=3, stride=strides[1], padding=1, bias=False),
			nn.BatchNorm2d(1024),
			nn.ReLU(inplace=True))
		
		self.conv_down3 = nn.Sequential(
			nn.Conv2d(1024, 2048, kernel_size=3, stride=strides[2], padding=1, bias=False),
			nn.BatchNorm2d(2048),
			nn.ReLU(inplace=True))
		
		self.decoder = build_decoder(num_classes, self.backone, BatchNorm)
		self.avgpool = nn.AvgPool2d(8, stride=1)
		self.fc = nn.Linear(2048*2, cls_classes)
		self.sigmoid = nn.Sigmoid()


		self.gamma1_1 = nn.Parameter(torch.zeros(1))
		self.gamma1_2 = nn.Parameter(torch.zeros(1))
		self.gamma2_1 = nn.Parameter(torch.zeros(1))
		self.gamma2_2 = nn.Parameter(torch.zeros(1))
		self.gamma2_3 = nn.Parameter(torch.zeros(1))
		self.gamma3_1 = nn.Parameter(torch.zeros(1))
		self.gamma3_2 = nn.Parameter(torch.zeros(1))
		self.gamma3_3 = nn.Parameter(torch.zeros(1))
		self.gamma4_1 = nn.Parameter(torch.zeros(1))
		self.gamma4_2 = nn.Parameter(torch.zeros(1))
		self.gamma4_3 = nn.Parameter(torch.zeros(1))
		
	def forward(self, x):
		c_x = self.backone_cls.conv1(x)
		c_x = self.backone_cls.bn1(c_x)
		c_x = self.backone_cls.relu(c_x)
		c_x = self.backone_cls.maxpool(c_x)
		
		s_x = self.backone_seg.conv1(x)
		s_x = self.backone_seg.bn1(s_x)
		s_x = self.backone_seg.relu(s_x)
		s_x = self.backone_seg.maxpool(s_x)
		
		c_x = self.backone_cls.layer1(c_x)
		s_x = self.backone_seg.layer1(s_x)  # 256
		low = s_x
		cs = self.gamma1_1 * c_x + self.gamma1_2 * s_x  # 256
		
		c_x = self.backone_cls.layer2(c_x)
		s_x = self.backone_seg.layer2(s_x)  # 512
		
		cs = self.conv_down1(cs)  # 512
		cs = self.gamma2_1 * cs + self.gamma2_2 * c_x + self.gamma2_3 * s_x
		
		c_x = self.backone_cls.layer3(c_x)
		s_x = self.backone_seg.layer3(s_x)  # 1024
		cs = self.conv_down2(cs)  # 1024
		cs = self.gamma3_1 * cs + self.gamma3_2 * c_x + self.gamma3_3 * s_x
		
		c_x = self.backone_cls.layer4(c_x)
		s_x = self.backone_seg.layer4(s_x)  # 2048
		cs = self.conv_down3(cs)  # 2048
		cs = self.gamma4_1 * cs + self.gamma4_2 * c_x + self.gamma4_3 * s_x
		
		# 分类部分：
		c_x = self.backone_cls.avgpool1(c_x)
		c_x = c_x.view(c_x.size(0), -1)
		
		cs_c = self.avgpool(cs)
		cs_c = cs_c.view(cs_c.size(0), -1)

		c_x = torch.cat([c_x, cs_c], dim=1)
		c_x = self.fc(c_x)
		
		# 分各部分：拼接 融合特征+分割特征，解码，上采样
		s_x = self.aspp_s(s_x)
		cs_s = self.conv_cs(cs)
		s_x = torch.cat([s_x, cs_s], dim=1)
		s_x = self.conv_s_cs(s_x)
		
		s_x = self.decoder(s_x, low)
		s_x = F.interpolate(s_x, size=x.size()[2:], mode='bilinear', align_corners=True)
		
		# print(s_x.size())
		# print(c_x.size())
		s_x = self.sigmoid(s_x)
		
		return s_x, c_x

BaseNet_version = 'SC_V1'

if __name__ == '__main__':
	net = Network(output_stride=32, num_classes=3, sync_bn=True, cls_classes=3)
	net.cuda()
	torchsummary.summary(net, (3, 256, 256))